﻿#include <limits>
#define _USE_MATH_DEFINES
#include <cmath>
#include <iostream>
#include <fstream>
#include <vector>
#include <algorithm>
#include <memory>

#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"
#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "stb_image_write.h"
#define STB_TRUETYPE_IMPLEMENTATION
#include "stb_truetype.h"

#include "geometry.h"
#include "model.h"

float fraction(float u) {
	while (u < 0) u += 1;
	while (u > 1) u -= 1;
	return u;
}

template <size_t Channel> struct _Texture2D {
	std::shared_ptr<unsigned char> tex;
	int width, height;
	_Texture2D(std::shared_ptr<unsigned char> t, int w, int h) : tex(t), width(w), height(h) {};
	vec<Channel, float> Sample(float u, float v) const {
		if (tex == nullptr) return vec<Channel, float>();
		int x = fraction(u) * (width - 1);
		int y = fraction(v) * (height - 1);
		vec<Channel, float> ret;
		for (int i = 0; i < Channel; i++) {
			int pixelPos = (y * width + x) * Channel + i;
			unsigned char cc = tex.get()[pixelPos];
			ret[i] = cc / 255.0f;
		}
		return ret;
	};
};

typedef _Texture2D<3> Texture2D;


void outputImg(const char* outpath, std::vector<Vec3f>& framebuffer, int width, int height, int channel) {
	char* outImg = (char*)std::malloc(width * height * channel);
	for (int i = 0; i < height * width; i++) {
		Vec3f& c = framebuffer[i];
		float cmax = std::max(c[0], std::max(c[1], c[2]));
		if (cmax > 1) c = c * (1. / cmax);
		for (int j = 0; j < channel; j++) {
			outImg[i * channel + j] = (char)(255 * std::max(0.f, std::min(1.f, c[j])));
		}
	}
	stbi_write_png(outpath, width, height, channel, outImg, 0);
	std::free(outImg);
}

struct Light {
	Vec3f position;
	float intensity;
	Light(const Vec3f& p, const float& i) : position(p), intensity(i) {};
};

struct Material {
	float refractive_index;
	Vec4f albedo;
	Vec3f diffuse_color;
	float specular_exponent;
	Material(const float& r, const Vec4f& a, const Vec3f& color, const float& spec) : refractive_index(r), albedo(a), diffuse_color(color), specular_exponent(spec) {}
	Material() : refractive_index(1), albedo(1, 0, 0, 0), diffuse_color(), specular_exponent() {}
};

struct Sphere {
	Vec3f center;
	float radius;
	Material material;

	Sphere(const Vec3f& c, const float& r, const Material& m) : center(c), radius(r), material(m) {}

	bool ray_intersect(const Vec3f& orig, const Vec3f& dir, float& t0) const {
		Vec3f L = center - orig;
		float tca = L * dir;
		float d2 = L * L - tca * tca;
		if (d2 > radius * radius) return false;
		float thc = sqrtf(radius * radius - d2);
		t0 = tca - thc;
		float t1 = tca + thc;
		if (t0 < 0) t0 = t1;
		if (t0 < 0) return false;
		return true;
	}

};

// https://en.wikipedia.org/wiki/Equirectangular_projection
struct Skybox {
	Texture2D tex;
	Skybox(const Texture2D& tex2D) : tex(tex2D) {};
	Vec3f ray_intersect(const Vec3f& dir) const {
		float longitude = atan2(-dir.z, dir.x) / M_PI;
		float latitude = acos(dir.y) / M_PI;
		return tex.Sample(longitude, latitude);
	};
};

struct Mesh {
	Model model;
	Material material;
	Mesh(const Model& mod, const Material& m) : model(mod), material(m) {}
	bool ray_triangle_intersect(const int& fi, const Vec3f& orig, const Vec3f& dir, float& near) const {
		return model.ray_triangle_intersect(fi, orig, dir, near);
	}
};

Vec3f reflect(const Vec3f& I, const Vec3f& N) {
	return I - N * (I * N) * 2.0f;
}

// https://graphics.stanford.edu/courses/cs148-10-summer/docs/2006--degreve--reflection_refraction.pdf
Vec3f refract(const Vec3f& I, const Vec3f& N, const float& refractive_index) {
	float cosi = -std::max(-1.0f, std::min(1.0f, (float)(I * N)));
	float etai = 1, etat = refractive_index;
	Vec3f n = N;
	if (cosi < 0) { // if the ray is inside the object, swap the indices and invert the normal to get the correct result
		cosi = -cosi;
		std::swap(etai, etat);
		n = -N;
	}
	float eta = etai / etat;
	float k = 1 - eta * eta * (1 - cosi * cosi);
	return k < 0 ? Vec3f(0, 0, 0) : I * eta + n * (eta * cosi - sqrt(k));
}

bool scene_intersect(const Vec3f& orig, const Vec3f& dir, const std::vector<Sphere>& spheres, const Mesh& mesh, Vec3f& hit, Vec3f& N, Material& material) {
	float spheres_dist = std::numeric_limits<float>::max();
	for (size_t i = 0; i < spheres.size(); i++) {
		float dist_i;
		if (spheres[i].ray_intersect(orig, dir, dist_i) && dist_i < spheres_dist) {
			spheres_dist = dist_i;
			hit = orig + dir * dist_i;
			N = (hit - spheres[i].center).normalize();
			material = spheres[i].material;
		}
	}

	const Model& model = mesh.model;
	Vec3f min, max;
	model.get_bbox(min, max);
	//Vec3f boundHit;
	if (HitBoundingBoxP(min, max, orig, dir)) {
		for (size_t i = 0; i < model.nfaces(); i++) {
			float dist_i;
			if (model.ray_triangle_intersect(i, orig, dir, dist_i) && dist_i < spheres_dist) {
				spheres_dist = dist_i;
				hit = orig + dir * dist_i;
				Vec3f edge1 = model.point(model.vert(i, 1)) - model.point(model.vert(i, 0));
				Vec3f edge2 = model.point(model.vert(i, 2)) - model.point(model.vert(i, 0));
				N = cross(edge1, edge2);
				material = mesh.material;
			}
		}
	}

	float checkboard_dist = std::numeric_limits<float>::max();
	if (fabs(dir.y) > 1e-3) {
		float d = -(orig.y + 4) / dir.y; // the checkerboard plane has equation y = -4
		Vec3f pt = orig + dir * d;
		if (d > 0 && fabs(pt.x) < 10 && pt.z<-10 && pt.z>-30 && d < spheres_dist) {
			checkboard_dist = d;
			hit = pt;
			N = Vec3f(0, 1, 0);
			material.diffuse_color = (int(0.5f * hit.x + 1000) + int(0.5f * hit.z)) & 1 ? Vec3f(1, 1, 1) : Vec3f(1, 0.7, 0.3);
			material.diffuse_color = material.diffuse_color * 0.3f;
		}
	}
	return std::min(spheres_dist, checkboard_dist) < 1000;
}

Vec3f cast_ray(const Vec3f& orig, const Vec3f& dir, const std::vector<Sphere>& spheres, const std::vector<Light>& lights, const Skybox& skybox, const Mesh& mesh, size_t depth = 0) {
	Vec3f point, N;
	Material material;
	if (depth > 4 || !scene_intersect(orig, dir, spheres, mesh, point, N, material)) {
		return skybox.ray_intersect(dir);
	}
	Vec3f reflect_dir = reflect(dir, N).normalize();
	Vec3f refract_dir = refract(dir, N, material.refractive_index).normalize();
	Vec3f reflect_orig = reflect_dir * N < 0 ? point - N * 1e-3 : point + N * 1e-3; // offset the original point to avoid occlusion by the object itself
	Vec3f refract_orig = refract_dir * N < 0 ? point - N * 1e-3 : point + N * 1e-3;
	Vec3f reflect_color = cast_ray(reflect_orig, reflect_dir, spheres, lights, skybox, mesh, depth + 1);
	Vec3f refract_color = cast_ray(refract_orig, refract_dir, spheres, lights, skybox, mesh, depth + 1);

	float diffuse_light_intensity = 0, specular_light_intensity = 0;
	for (size_t i = 0; i < lights.size(); i++) {
		Vec3f light_dir = (lights[i].position - point).normalize();
		float light_distance = (lights[i].position - point).norm();

		Vec3f shadow_orig = light_dir * N < 0 ? point - N * 1e-3 : point + N * 1e-3;
		Vec3f shadow_pt, shadow_N;
		Material tepmaterial;
		if (scene_intersect(shadow_orig, light_dir, spheres, mesh, shadow_pt, shadow_N, tepmaterial) && (shadow_pt - shadow_orig).norm() < light_distance)
			continue;

		diffuse_light_intensity += lights[i].intensity * std::max((float)(light_dir * N), 0.0f);
		specular_light_intensity += powf(std::max(0.0f, (float)(reflect(light_dir, N) * dir)), material.specular_exponent) * lights[i].intensity;
	}
	return material.diffuse_color * diffuse_light_intensity * material.albedo[0] +
		Vec3f(1.f, 1.f, 1.f) * specular_light_intensity * material.albedo[1] +
		reflect_color * material.albedo[2] + refract_color * material.albedo[3];
}

void render(const std::vector<Sphere>& spheres, const std::vector<Light>& lights, const Skybox& skybox, const Mesh& mesh) {
	const int width = 1024;
	const int height = 768;
	const float fov = M_PI / 2.0f;

	std::vector<Vec3f> framebuffer(width * height);

#pragma omp parallel for
	for (size_t j = 0; j < height; j++) {
		for (size_t i = 0; i < width; i++) {
			float x = (2 * (i + 0.5f) / (float)width - 1) * tan(fov / 2.) * width / (float)height;
			float y = -(2 * (j + 0.5f) / (float)height - 1) * tan(fov / 2.);
			Vec3f dir = Vec3f(x, y, -1).normalize();
			framebuffer[i + j * width] = cast_ray(Vec3f(0, 0, 0), dir, spheres, lights, skybox, mesh);
		}
	}
	outputImg("./TestImg.png", framebuffer, width, height, 3);
}

int main() {
	Material      ivory(1.0, Vec4f(0.6, 0.3, 0.1, 0.0), Vec3f(0.4, 0.4, 0.3), 50.);
	Material      glass(1.5, Vec4f(0.0, 0.5, 0.1, 0.8), Vec3f(0.6, 0.7, 0.8), 125.);
	Material red_rubber(1.0, Vec4f(0.9, 0.1, 0.0, 0.0), Vec3f(0.3, 0.1, 0.1), 10.);
	Material     mirror(1.0, Vec4f(0.0, 10.0, 0.8, 0.0), Vec3f(1.0, 1.0, 1.0), 1425.);
	Material      glassduck(1.5, Vec4f(0.3, 0.5, 0.1, 0.8), Vec3f(0.6, 0.7, 0.8), 125.);

	std::vector<Sphere> spheres;
	spheres.push_back(Sphere(Vec3f(-3, 0, -16), 2, ivory));
	spheres.push_back(Sphere(Vec3f(-1.0, -1.5, -12), 2, glass));
	spheres.push_back(Sphere(Vec3f(1.5, -0.5, -18), 3, red_rubber));
	spheres.push_back(Sphere(Vec3f(7, 5, -18), 4, mirror));

	std::vector<Light> lights;
	lights.push_back(Light(Vec3f(-20, 20, 20), 1.5f));
	lights.push_back(Light(Vec3f(30, 50, -25), 1.8));
	lights.push_back(Light(Vec3f(30, 20, 30), 1.7));

	int x, y, n;
	unsigned char* texdata = stbi_load("./envmap.jpg", &x, &y, &n, 0);
	std::shared_ptr<unsigned char> skyTexPtr(texdata);
	Texture2D skyTex(skyTexPtr, x, y);
	Skybox skybox(skyTex);

	Model duckModel("./duck.obj");
	Mesh duckMesh(duckModel, glassduck);

	render(spheres, lights, skybox, duckMesh);
	return 1;
}
